/**
* @file post_process.h
*
* Copyright (C) 2020. Huawei Technologies Co., Ltd. All rights reserved.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
*/
#pragma once

#include <iostream>
#include <vector>
#include "utils.h"
#include <opencv2/opencv.hpp>
using namespace std;



class PostProcess {
public:
    /**
    * @brief Constructor
    * @param [in] outputClass: ouput buffer for class
    * @param [in] outputBox: ouput buffer for box
    * @param [in] xScale: ouput buffer for x scale
    * @param [in] yScale: ouput buffer for y scale
    * @param [in] originImage: origin image path
    */
    PostProcess(float *outputClass, uint32_t modelWidth, uint32_t modelHeight) :
            outputClass_(outputClass), modelWidth_(modelWidth),  modelHeight_(modelHeight){}

    /**
    * @brief Constructor
    */
    PostProcess() : outputClass_(nullptr) {}

    /**
    * @brief Destructor
    */
    virtual ~PostProcess() {}

    /**
    * @brief post process
    * @return result
    */
    Result Process(cv::Mat& image, vector<BBox>& result);

   /**
    * @brief draw detection boxes
    */
    static void DrawDetections(cv::Mat& ori_img, std::vector<BBox>& detections);


private:
    void SetBoxInfo(size_t index, BBox &box);

    void NMS(vector<BBox> &boxes, const float threshold);

    float IOU(const BBox &b1, const BBox &b2);

    void DrawBoundBoxToImage_ycf(const std::vector <BBox> &result, cv::Mat& imageInput);

    float *outputClass_;
    float xScale_;
    float yScale_;

    uint32_t modelWidth_;      // 模型输入的宽度
    uint32_t modelHeight_;     // 高度

    uint32_t imageWidth_;      // 后处理的目标图片宽度高度
    uint32_t imageHeight_; 
};

